package com.omnisyr.textrebuilder.util.structure;

import java.util.Comparator;

public class SimpleTreeMap<K, V> {

	private Node root;
	private Comparator<K> comparator;

	public SimpleTreeMap(Comparator<K> comparator) {

		this.comparator = comparator;

	}

	private class Node {

		private K key;
		private V value;
		private Node left;
		private Node right;
		private boolean isRed;

		private Node(K key, V value, boolean isRed) {

			this.key = key;
			this.value = value;
			this.isRed = isRed;

		}

	}

	public void add(K key, V value) {

		this.root = this.add(this.root, key, value);
		this.root.isRed = false;

	}

	private Node add(Node node, K key, V value) {

		if (node == null) {

			return new Node(key, value, true);

		}

		int compareResult = this.comparator.compare(key, node.key);

		if (compareResult < 0) {

			node.left = this.add(node.left, key, value);

		} else if (compareResult > 0) {

			node.right = this.add(node.right, key, value);

		} else {

			node.value = value;

			return node;

		}

		return this.ensureBalance(node);

	}

	public V get(K key) {

		Node node = this.root;

		while (node != null) {

			int compareResult = this.comparator.compare(key, node.key);

			if (compareResult < 0) {

				node = node.left;

				continue;

			}

			if (compareResult > 0) {

				node = node.right;

				continue;

			}

			return node.value;

		}

		return null;

	}

	public void delete(K key) {

		if (this.get(key) == null) {

			return;

		}

		if (!this.isRed(root.left) && !this.isRed(root.right)) {

			this.root.isRed = true;

		}

		this.root = this.delete(this.root, key);

		if (!isEmpty()) {

			this.root.isRed = false;

		}

	}

	private Node delete(Node node, K key) {

		if (this.comparator.compare(key, node.key) < 0) {

			if (!this.isRed(node.left) && !this.isRed(node.left.left)) {

				node = this.moveRedLeft(node);

			}

			node.left = this.delete(node.left, key);

		} else {

			if (isRed(node.left)) {

				node = this.rotateRight(node);

			}

			if (this.comparator.compare(key, node.key) == 0 && (node.right == null)) {

				return null;

			}

			if (!isRed(node.right) && !isRed(node.right.left)) {

				node = this.moveRedRight(node);

			}

			if (this.comparator.compare(key, node.key) == 0) {

				Node minNode = this.min(node.right);

				node.key = minNode.key;
				node.value = minNode.value;
				node.right = this.deleteMin(node.right);

			} else {

				node.right = delete(node.right, key);

			}

		}

		return this.ensureBalance(node);

	}

	public boolean isEmpty() {

		return this.root == null;

	}

	private boolean isRed(Node node) {

		if (node == null) {

			return false;

		}

		return node.isRed;

	}

	private Node moveRedLeft(Node node) {

		this.flipColors(node);

		if (this.isRed(node.right.left)) {

			node.right = this.rotateRight(node.right);

			node = this.rotateLeft(node);

			this.flipColors(node);

		}

		return node;

	}

	private Node moveRedRight(Node node) {

		this.flipColors(node);

		if (this.isRed(node.left.left)) {

			node = this.rotateRight(node);

			this.flipColors(node);

		}

		return node;

	}

	private Node min(Node node) {

		if (node.left == null) {

			return node;

		}

		return this.min(node.left);

	}

	private Node deleteMin(Node node) {

		if (node.left == null) {

			return null;

		}

		if (!this.isRed(node.left) && !this.isRed(node.left.left)) {

			node = this.moveRedLeft(node);

		}

		node.left = this.deleteMin(node.left);

		return this.ensureBalance(node);

	}

	private Node ensureBalance(Node node) {

		if (this.isRed(node.right) && !this.isRed(node.left)) {

			node = this.rotateLeft(node);

		}

		if (this.isRed(node.left) && this.isRed(node.left.left)) {

			node = this.rotateRight(node);

		}

		if (this.isRed(node.left) && this.isRed(node.right)) {

			this.flipColors(node);

		}

		return node;

	}

	private Node rotateLeft(Node node) {

		Node rightNode = node.right;

		node.right = rightNode.left;
		rightNode.left = node;
		rightNode.isRed = node.isRed;
		node.isRed = true;

		return rightNode;

	}

	private Node rotateRight(Node node) {

		Node leftNode = node.left;

		node.left = leftNode.right;
		leftNode.right = node;
		leftNode.isRed = node.isRed;
		node.isRed = true;

		return leftNode;

	}

	private void flipColors(Node node) {

		node.isRed = true;
		node.left.isRed = false;
		node.right.isRed = false;

	}

}
